import Data.Char (isDigit)
import Utils (readInt)

type Board     = [String]
type Coord     = (Int, Int)
type Symbol    = Coord
type Number    = (Coord, Int)
type NumberBox = (Coord, Coord)

type Star      = Symbol

getNumbers :: Board -> [Number]
getNumbers = 
    go 0 0
    where
        go :: Int -> Int -> Board -> [Number]
        go _ _ []         = []
        go r _ ([]:lists) = go (r + 1) 0 lists
        go r c (list@(x:xs):lists)
            | (not . isDigit) x = go r (c + 1) (xs:lists)
            | otherwise         = ((r, c), number) : go r c' (xs':lists)
            where
                number = readInt . takeWhile isDigit $ list
                offset = length . show $ number
                xs'    = drop offset list
                c'     = c + offset

box :: Number -> NumberBox
box ((r, c), number) = 
    ((r - 1, c - 1), (r + 1, c2))
    where
        c2 = c + (length . show $ number)


------------
-- Part 1 --
------------

getSymbols :: Board -> [Symbol]
getSymbols = 
    go 0 0
    where
        go _ _ []         = []
        go r _ ([]:lists) = go (r + 1) 0 lists
        go r c ((x:xs):lists)
            | (isDigit x) || (x == '.') = go r (c + 1) (xs:lists)
            | otherwise                 = (r, c) : go r (c + 1) (xs:lists)

isPartNb :: [Symbol] -> Number -> Bool
isPartNb symbols number = any (within numberox) symbols
    where
        numberox                           = box number
        within ((r1, c1), (r2, c2)) (r, c) = (r >= r1) && (r <= r2) && (c >= c1) && (c <= c2)

------------
-- Part 2 --
------------

getStars :: Board -> [Star]
getStars = 
    go 0 0
    where
        go _ _ []             = []
        go r _ ([]:lists)     = go (r + 1) 0 lists
        go r c ((x:xs):lists)
            | x == '*'  = (r, c) : go r (c + 1) (xs:lists)
            | otherwise = go r (c + 1) (xs:lists)

numberAdjToStar :: Star -> Number -> Bool
numberAdjToStar (r, c) number = 
    (r >= r1) && (r <= r2) && (c >= c1) && (c <= c2)
    where
        ((r1, c2), (r2, c2)) = box number

gearRatio :: [Number] -> Star -> Int
gearRatio numbers star = 
    case (length adjacentNumbers) of
        2 -> product $ map snd $ adjacentNumbers
        _ -> 0
    where
        adjacentNumbers = filter (numberAdjToStar star) numbers

main = do
    contents <- readFile "input.txt"

    let board = lines contents
    let symbols = getSymbols board
    let numbers = getNumbers board
    
    print $  sum . map snd . filter (isPartNb symbols) $ number

    let stars = getStars board
    print $ sum . map (gearRatio numbers) $ stars